import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt

# ---------- 参数 ----------
T = 10.0                      # 轨迹周期 (秒)
dt = 0.004                   # 时间间隔
N_harmonics = 3             # 五阶Fourier
omega = 2 * np.pi / T       # 基频
t = np.arange(0, T, dt)     # 时间向量

# ---------- Fourier 基函数 ----------
def fourier_series(t, coeffs):
    a0 = coeffs[0]
    result = a0 * np.ones_like(t)
    for k in range(1, N_harmonics+1):
        a_k = coeffs[2*k - 1]
        b_k = coeffs[2*k]
        result += a_k * np.cos(k * omega * t) + b_k * np.sin(k * omega * t)
    return result

def d_fourier_series(t, coeffs):  # 一阶导
    result = np.zeros_like(t)
    for k in range(1, N_harmonics+1):
        a_k = coeffs[2*k - 1]
        b_k = coeffs[2*k]
        result += -k * omega * a_k * np.sin(k * omega * t) + k * omega * b_k * np.cos(k * omega * t)
    return result

def dd_fourier_series(t, coeffs):  # 二阶导
    result = np.zeros_like(t)
    for k in range(1, N_harmonics+1):
        a_k = coeffs[2*k - 1]
        b_k = coeffs[2*k]
        result += -(k * omega)**2 * a_k * np.cos(k * omega * t) - (k * omega)**2 * b_k * np.sin(k * omega * t)
    return result

# ---------- 目标函数 ----------
def objective(coeffs):
    x = fourier_series(t, coeffs)
    return -np.std(x)  # 最大化轨迹活跃程度

# ---------- 约束 ----------
def constraint_start_end_velocity(coeffs):
    v = d_fourier_series(np.array([0.0, T]), coeffs)
    return v

def constraint_start_end_acceleration(coeffs):
    a = dd_fourier_series(np.array([0.0, T]), coeffs)
    return a

def constraint_position_bounds(coeffs):
    x = fourier_series(t, coeffs)
    return np.array([3 - np.max(np.abs(x))])

def constraint_velocity_bounds(coeffs):
    v = d_fourier_series(t, coeffs)
    return np.array([0.001 - np.max(np.abs(v))])

def constraint_acceleration_bounds(coeffs):
    a = dd_fourier_series(t, coeffs)
    return np.array([0.01 - np.max(np.abs(a))])

# ---------- 初始猜测 ----------
x0 = np.zeros(2 * N_harmonics + 1)  # a0 + 5*(a_k, b_k)

# ---------- 构造约束列表 ----------
constraints = [
    {'type': 'eq', 'fun': constraint_start_end_velocity},
    {'type': 'eq', 'fun': constraint_start_end_acceleration},
    {'type': 'ineq', 'fun': constraint_position_bounds},
    {'type': 'ineq', 'fun': constraint_velocity_bounds},
    {'type': 'ineq', 'fun': constraint_acceleration_bounds},
]

# ---------- 优化求解 ----------
trajectory = []

for i in range(7):
    x0 = np.random.uniform(-0.5, 0.5, 2 * N_harmonics + 1)
    res = minimize(objective, x0, method='SLSQP', constraints=constraints)
    # ---------- 绘图 ----------
    x = fourier_series(t, res.x)
    v = d_fourier_series(t, res.x)
    a = dd_fourier_series(t, res.x)

    trajectory.append(x)

np.savetxt("seven_fourier_traj.txt", list(map(list, zip(*trajectory))), fmt="%.5f", delimiter=',')

# plt.figure(figsize=(10,6))
# plt.subplot(3,1,1); plt.plot(t, x); plt.title('Position'); plt.grid(); plt.ylim(-3.5, 3.5)
# plt.subplot(3,1,2); plt.plot(t, v); plt.title('Velocity'); plt.grid(); plt.ylim(-1.5, 1.5)
# plt.subplot(3,1,3); plt.plot(t, a); plt.title('Acceleration'); plt.grid(); plt.ylim(-6, 6)
# plt.tight_layout()
# plt.show()
